"""
音频-视觉事件数据集模块
该模块定义了AVEDataset类，用于加载和处理音频-视觉事件数据集
使用HDF5格式存储特征数据和标签
"""

# 导入操作系统接口模块，提供文件路径操作功能
import os
# 导入HDF5文件操作库，用于处理大型科学数据集
import h5py
# 导入PyTorch深度学习框架
import torch
# 从PyTorch数据工具包中导入数据集和数据加载器相关类
from torch.utils.data import Dataset, DataLoader

# 定义音频-视觉事件数据集类，继承自PyTorch的Dataset类
class AVEDataset(Dataset):
    """
    音频-视觉事件数据集类
    用于加载音频和视觉特征数据及其对应的标签
    支持训练集、验证集和测试集的分割
    """
    
    def __init__(self, data_root, split='train'):
        """
        初始化数据集
        
        参数:
        data_root: 数据根目录路径，包含所有特征文件和标签文件
        split: 数据集分割类型，默认为'train'，可选'train'、'val'、'test'等
        """
        
        # 调用父类Dataset的初始化方法
        super(AVEDataset, self).__init__()
        # 设置数据集分割类型（训练集、验证集或测试集）
        self.split = split
        # 构建视觉特征文件的完整路径
        self.visual_feature_path = os.path.join(data_root, 'visual_feature.h5')
        # 构建音频特征文件的完整路径
        self.audio_feature_path = os.path.join(data_root, 'audio_feature.h5')
        
        # 注释说明：以下用于监督学习任务
        # 构建标签文件的完整路径
        self.labels_path = os.path.join(data_root, 'labels.h5')
        # 构建样本顺序文件的完整路径，文件名根据分割类型动态生成
        self.sample_order_path = os.path.join(data_root, f'{split}_order.h5')
        # 标记HDF5文件是否已打开，初始状态为未打开
        self.h5_isOpen = False

    def __getitem__(self, index):
        """
        获取指定索引的数据样本
        
        参数:
        index: 样本索引，整数类型
        
        返回:
        visual_feat: 视觉特征，numpy数组或tensor
        audio_feat: 音频特征，numpy数组或tensor  
        label: 样本标签，numpy数组或tensor
        """
        
        # 如果HDF5文件尚未打开，则打开所有需要的HDF5文件
        if not self.h5_isOpen:
            # 打开视觉特征HDF5文件并获取数据集
            self.visual_feature = h5py.File(self.visual_feature_path, 'r')['avadataset']
            # 打开音频特征HDF5文件并获取数据集
            self.audio_feature = h5py.File(self.audio_feature_path, 'r')['avadataset']
            # 打开标签HDF5文件并获取数据集
            self.labels = h5py.File(self.labels_path, 'r')['avadataset']
            # 打开样本顺序HDF5文件并获取顺序数组
            self.sample_order = h5py.File(self.sample_order_path, 'r')['order']
            # 标记HDF5文件已打开，避免重复打开
            self.h5_isOpen = True
        
        # 根据索引从样本顺序数组中获取实际的样本索引
        sample_index = self.sample_order[index]
        # 从视觉特征数据集中获取对应样本的视觉特征
        visual_feat = self.visual_feature[sample_index]
        # 从音频特征数据集中获取对应样本的音频特征
        audio_feat = self.audio_feature[sample_index]
        # 从标签数据集中获取对应样本的标签
        label = self.labels[sample_index]

        # 返回样本的视觉特征、音频特征和标签
        return visual_feat, audio_feat, label

    def __len__(self):
        """
        获取数据集的样本数量
        
        返回:
        sample_num: 数据集的样本总数，整数类型
        """
        
        # 打开样本顺序HDF5文件
        f = h5py.File(self.sample_order_path, 'r')
        # 获取顺序数组的长度，即样本总数
        sample_num = len(f['order'])
        # 关闭HDF5文件
        f.close()
        # 返回样本数量
        return sample_num